import time
import threading
from typing import Optional

class SnowflakeIdGenerator:
    """
    雪花算法ID生成器
    
    雪花算法生成的ID结构：
    - 1位符号位（固定为0）
    - 41位时间戳（毫秒级，可使用69年）
    - 10位机器ID（5位数据中心ID + 5位机器ID）
    - 12位序列号（同一毫秒内的序列号）
    
    总共64位，生成的ID为正整数
    """
    
    def __init__(self, datacenter_id: int = 1, worker_id: int = 1):
        """
        初始化雪花算法生成器
        
        Args:
            datacenter_id: 数据中心ID (0-31)
            worker_id: 机器ID (0-31)
        """
        # 各部分位数
        self.WORKER_ID_BITS = 5
        self.DATACENTER_ID_BITS = 5
        self.SEQUENCE_BITS = 12
        
        # 最大值
        self.MAX_WORKER_ID = (1 << self.WORKER_ID_BITS) - 1
        self.MAX_DATACENTER_ID = (1 << self.DATACENTER_ID_BITS) - 1
        self.MAX_SEQUENCE = (1 << self.SEQUENCE_BITS) - 1
        
        # 位移量
        self.WORKER_ID_SHIFT = self.SEQUENCE_BITS
        self.DATACENTER_ID_SHIFT = self.SEQUENCE_BITS + self.WORKER_ID_BITS
        self.TIMESTAMP_LEFT_SHIFT = self.SEQUENCE_BITS + self.WORKER_ID_BITS + self.DATACENTER_ID_BITS
        
        # 起始时间戳 (2020-01-01 00:00:00)
        self.EPOCH = 1577836800000
        
        # 验证参数
        if datacenter_id > self.MAX_DATACENTER_ID or datacenter_id < 0:
            raise ValueError(f"datacenter_id must be between 0 and {self.MAX_DATACENTER_ID}")
        if worker_id > self.MAX_WORKER_ID or worker_id < 0:
            raise ValueError(f"worker_id must be between 0 and {self.MAX_WORKER_ID}")
        
        self.datacenter_id = datacenter_id
        self.worker_id = worker_id
        self.sequence = 0
        self.last_timestamp = -1
        self.lock = threading.Lock()
    
    def _current_millis(self) -> int:
        """获取当前时间戳（毫秒）"""
        return int(time.time() * 1000)
    
    def _wait_next_millis(self, last_timestamp: int) -> int:
        """等待下一毫秒"""
        timestamp = self._current_millis()
        while timestamp <= last_timestamp:
            timestamp = self._current_millis()
        return timestamp
    
    def next_id(self) -> int:
        """
        生成下一个ID
        
        Returns:
            生成的唯一ID
        """
        with self.lock:
            timestamp = self._current_millis()
            
            # 时钟回拨检查
            if timestamp < self.last_timestamp:
                raise RuntimeError(f"Clock moved backwards. Refusing to generate id for {self.last_timestamp - timestamp} milliseconds")
            
            # 同一毫秒内
            if timestamp == self.last_timestamp:
                self.sequence = (self.sequence + 1) & self.MAX_SEQUENCE
                # 序列号溢出，等待下一毫秒
                if self.sequence == 0:
                    timestamp = self._wait_next_millis(self.last_timestamp)
            else:
                # 新的毫秒，序列号重置
                self.sequence = 0
            
            self.last_timestamp = timestamp
            
            # 生成ID
            return ((timestamp - self.EPOCH) << self.TIMESTAMP_LEFT_SHIFT) | \
                   (self.datacenter_id << self.DATACENTER_ID_SHIFT) | \
                   (self.worker_id << self.WORKER_ID_SHIFT) | \
                   self.sequence
    
    def parse_id(self, snowflake_id: int) -> dict:
        """
        解析雪花算法ID
        
        Args:
            snowflake_id: 雪花算法生成的ID
            
        Returns:
            包含时间戳、数据中心ID、机器ID、序列号的字典
        """
        timestamp = ((snowflake_id >> self.TIMESTAMP_LEFT_SHIFT) + self.EPOCH)
        datacenter_id = (snowflake_id >> self.DATACENTER_ID_SHIFT) & ((1 << self.DATACENTER_ID_BITS) - 1)
        worker_id = (snowflake_id >> self.WORKER_ID_SHIFT) & ((1 << self.WORKER_ID_BITS) - 1)
        sequence = snowflake_id & ((1 << self.SEQUENCE_BITS) - 1)
        
        return {
            'timestamp': timestamp,
            'datetime': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(timestamp / 1000)),
            'datacenter_id': datacenter_id,
            'worker_id': worker_id,
            'sequence': sequence
        }


# 全局雪花算法生成器实例
_snowflake_generator: Optional[SnowflakeIdGenerator] = None


def init_snowflake(datacenter_id: int = 1, worker_id: int = 1) -> None:
    """
    初始化全局雪花算法生成器
    
    Args:
        datacenter_id: 数据中心ID (0-31)
        worker_id: 机器ID (0-31)
    """
    global _snowflake_generator
    _snowflake_generator = SnowflakeIdGenerator(datacenter_id, worker_id)


def generate_id() -> int:
    """
    生成雪花算法ID
    
    Returns:
        生成的唯一ID
        
    Raises:
        RuntimeError: 如果雪花算法生成器未初始化
    """
    global _snowflake_generator
    if _snowflake_generator is None:
        # 使用默认配置初始化
        init_snowflake()
    return _snowflake_generator.next_id()


def parse_snowflake_id(snowflake_id: int) -> dict:
    """
    解析雪花算法ID
    
    Args:
        snowflake_id: 雪花算法生成的ID
        
    Returns:
        包含时间戳、数据中心ID、机器ID、序列号的字典
    """
    global _snowflake_generator
    if _snowflake_generator is None:
        init_snowflake()
    return _snowflake_generator.parse_id(snowflake_id)


# 为了向后兼容，提供一个简单的函数
def snowflake_id() -> int:
    """
    生成雪花算法ID（简化版本）
    
    Returns:
        生成的唯一ID
    """
    return generate_id()